// Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved.

// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
// WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#ifndef LSTM_FUSED_GATES_CALC_LAYER_H
#define LSTM_FUSED_GATES_CALC_LAYER_H

#include "training/base/layers/TrainableLayer.h"
#include "training/base/layers/parameters/trainable/LSTMFusedGatesCalcParams.h"
#include <training/base/common/Common.h>

namespace raul
{

class LSTMFusedGatesCalcLayer : public TrainableLayer
{

  public:
    LSTMFusedGatesCalcLayer(const Name& name, const LSTMFusedGatesCalcParams& params, NetworkParameters& networkParameters);

    LSTMFusedGatesCalcLayer(LSTMFusedGatesCalcLayer&&) = default;
    LSTMFusedGatesCalcLayer(const LSTMFusedGatesCalcLayer&) = delete;
    LSTMFusedGatesCalcLayer& operator=(const LSTMFusedGatesCalcLayer&) = delete;

  private:
    size_t mOutputsCount;
    bool mUseBias;
    dtype mZoneout;
    bool mUseSingleParamTensor;
    dtype mForgetBias;
    Name mTmpCalculationsName;
    Name mGatesName;

    // Needed if zoneout used
    Name mRandomNameHidden;
    Name mRandomNameCell;
    Name mNoZoneoutNewCellName;
    Name mNoZoneoutNewHiddenGradName;
    Name mNoZoneoutNewCellGradName;
    bool mUseZoneout;

    // Trainable names
    Names mTrainableParamsNames;

    template<typename MM>
    friend class LSTMFusedGatesCalcLayerCPU;
};

} // raul namespace

#endif